"""
均线上穿下穿算法可视化工具
功能：生成详细的可视化图表，展示趋势追踪止损中使用的均线上穿下穿算法
"""

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from tqsdk import TqApi, TqAuth, TqKq
from datetime import datetime

# MyTT模块中的函数已在下方直接定义，无需额外导入MyTT库
# 这些是从MyTT库中提取的常用技术指标函数
def REF(Series, N=1):
    """向前引用N周期的数据"""
    return pd.Series(Series).shift(N)

def EMA(Series, N=10):
    """指数移动平均线"""
    return pd.Series(Series).ewm(span=N, adjust=False).mean()

def SUM(Series, N=20):
    """求和"""
    return pd.Series(Series).rolling(N).sum()

def calculate_signals(klines):
    """
    计算交易信号
    :param klines: K线数据
    :return: 包含信号的DataFrame
    """
    # 转换为numpy数组以提高计算效率
    close = np.array(klines.close)
    low = np.array(klines.low)
    open_price = np.array(klines.open)
    high = np.array(klines.high)
    
    # 计算Q值 (与原始代码一致)
    Q = (3 * close + low + open_price + high) / 6
    
    # 计算trading_line (与原始代码一致)
    terms = [
        26 * Q,
        25 * REF(Q, 1),
        24 * REF(Q, 2),
        23 * REF(Q, 3),
        22 * REF(Q, 4),
        21 * REF(Q, 5),
        20 * REF(Q, 6),
        19 * REF(Q, 7),
        18 * REF(Q, 8),
        17 * REF(Q, 9),
        16 * REF(Q, 10),
        15 * REF(Q, 11),
        14 * REF(Q, 12),
        13 * REF(Q, 13),
        12 * REF(Q, 14),
        11 * REF(Q, 15),
        10 * REF(Q, 16),
        9 * REF(Q, 17),
        8 * REF(Q, 18),
        7 * REF(Q, 19),
        6 * REF(Q, 20),
        5 * REF(Q, 21),
        4 * REF(Q, 22),
        3 * REF(Q, 23),
        2 * REF(Q, 24),
        REF(Q, 25)
    ]
    
    trading_line = sum(terms) / 351
    short_line = EMA(trading_line, 7)
    
    # 创建结果DataFrame
    result = pd.DataFrame({
        'datetime': klines.datetime,
        'open': open_price,
        'high': high,
        'low': low,
        'close': close,
        'trading_line': trading_line,
        'short_line': short_line
    })
    
    # 转换时间戳为日期时间
    result['date'] = pd.to_datetime(result['datetime'], unit='s')
    
    # 计算信号
    result['signal'] = np.nan  # 默认为NaN
    result['signal_value'] = np.nan  # 用于绘图的信号值
    
    # 计算上穿下穿信号
    for i in range(2, len(result)):
        # 上穿信号: 前一周期trading_line < short_line，当前周期trading_line > short_line
        if (result.trading_line.iloc[i-1] < result.short_line.iloc[i-1] and 
            result.trading_line.iloc[i] > result.short_line.iloc[i]):
            result.loc[result.index[i], 'signal'] = '上穿'
            result.loc[result.index[i], 'signal_value'] = result.trading_line.iloc[i]
        
        # 下穿信号: 前一周期trading_line > short_line，当前周期trading_line < short_line
        elif (result.trading_line.iloc[i-1] > result.short_line.iloc[i-1] and 
              result.trading_line.iloc[i] < result.short_line.iloc[i]):
            result.loc[result.index[i], 'signal'] = '下穿'
            result.loc[result.index[i], 'signal_value'] = result.trading_line.iloc[i]
    
    return result

def create_visualization(symbol, period_name, result):
    """
    创建详细的可视化图表
    :param symbol: 合约代码
    :param period_name: 周期名称
    :param result: 计算结果DataFrame
    """
    # 创建一个大图表
    fig = plt.figure(figsize=(16, 12))
    
    # 设置标题
    fig.suptitle(f"{symbol} {period_name} 均线上穿下穿算法验证", fontsize=16)
    
    # 1. 绘制K线图和均线
    ax1 = plt.subplot2grid((3, 1), (0, 0), rowspan=2)
    
    # 绘制收盘价
    ax1.plot(result['date'], result['close'], label='收盘价', color='gray', alpha=0.5)
    
    # 绘制均线
    ax1.plot(result['date'], result['trading_line'], label='Trading Line', color='blue', linewidth=1.5)
    ax1.plot(result['date'], result['short_line'], label='Short Line', color='red', linewidth=1.5)
    
    # 标记上穿和下穿点
    up_cross = result[result['signal'] == '上穿']
    down_cross = result[result['signal'] == '下穿']
    
    ax1.scatter(up_cross['date'], up_cross['signal_value'], color='green', marker='^', s=100, label='上穿')
    ax1.scatter(down_cross['date'], down_cross['signal_value'], color='red', marker='v', s=100, label='下穿')
    
    # 添加网格和图例
    ax1.grid(True, alpha=0.3)
    ax1.legend(loc='upper left')
    ax1.set_title(f"{symbol} {period_name} K线和均线")
    ax1.set_ylabel('价格')
    
    # 格式化x轴日期
    ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ax1.xaxis.set_major_locator(mdates.AutoDateLocator())
    
    # 2. 绘制信号图
    ax2 = plt.subplot2grid((3, 1), (2, 0), rowspan=1, sharex=ax1)
    
    # 创建信号值列表 (1=上穿, -1=下穿, 0=无信号)
    signal_values = np.zeros(len(result))
    signal_values[result['signal'] == '上穿'] = 1
    signal_values[result['signal'] == '下穿'] = -1
    
    # 绘制信号柱状图
    ax2.bar(result['date'], signal_values, color=['green' if x == 1 else 'red' if x == -1 else 'gray' for x in signal_values], alpha=0.7)
    
    # 添加水平线
    ax2.axhline(y=0, color='black', linestyle='-', alpha=0.3)
    
    # 设置y轴刻度和标签
    ax2.set_yticks([-1, 0, 1])
    ax2.set_yticklabels(['下穿', '无信号', '上穿'])
    
    # 添加网格
    ax2.grid(True, alpha=0.3)
    ax2.set_title('信号分布')
    
    # 格式化x轴日期
    ax2.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
    ax2.xaxis.set_major_locator(mdates.AutoDateLocator())
    
    # 旋转x轴标签
    plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)
    plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.92)
    
    # 保存图表
    plt_filename = f"{symbol.replace('.', '_')}_{period_name}_详细图表.png"
    plt.savefig(plt_filename, dpi=150)
    print(f"详细图表已保存到文件: {plt_filename}")
    
    # 显示图表
    plt.show()

# 主函数
print("开始验证均线上穿下穿算法...")

try:
    # 创建API连接
    api = TqApi(account=TqKq(), auth=TqAuth("cps168", "alibaba"))
    
    # 设置默认合约和周期
    symbol = "SHFE.au2412"  # 上期所黄金
    period = 86400  # 日线
    period_name = "日线"
    
    print(f"正在获取 {symbol} 的{period_name}数据...")
    klines = api.get_kline_serial(symbol, period, data_length=500)
    
    # 计算信号
    result = calculate_signals(klines)
    
    # 输出结果摘要
    print("\n计算完成! 结果摘要:")
    print(f"总K线数量: {len(result)}")
    print(f"上穿信号数量: {result['signal'].value_counts().get('上穿', 0)}")
    print(f"下穿信号数量: {result['signal'].value_counts().get('下穿', 0)}")
    
    # 显示最近的10个信号
    recent_signals = result[result['signal'].notna()].tail(10)
    if not recent_signals.empty:
        print("\n最近的10个信号:")
        for idx, row in recent_signals.iterrows():
            date_str = row['date'].strftime('%Y-%m-%d %H:%M:%S') if hasattr(row['date'], 'strftime') else pd.to_datetime(row['datetime'], unit='s').strftime('%Y-%m-%d %H:%M:%S')
            print(f"{date_str}: {row['signal']} (trading_line: {row['trading_line']:.2f}, short_line: {row['short_line']:.2f})")
    
    # 保存结果到CSV
    csv_filename = f"{symbol.replace('.', '_')}_{period_name}_信号详情.csv"
    result.to_csv(csv_filename, index=False)
    print(f"\n结果已保存到文件: {csv_filename}")
    
    # 创建可视化图表
    create_visualization(symbol, period_name, result)
    
except Exception as e:
    print(f"发生错误: {e}")
finally:
    if 'api' in locals():
        api.close()
        print("API连接已关闭")